package com.lhf.springcloud.rule;

import com.netflix.loadbalancer.*;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.List;

/**
 * 自定义负载均衡规则
 */

public class SelectByInstanceRule extends RoundRobinRule {

    private ILoadBalancer lb;
//    private final static IRule DEFAULT_RULE = new RoundRobinRule(); // 可以维护一个对象，也可以使用继承的方式
    @Override
    public Server choose(Object key) {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        String instanceId = request.getParameter("instanceId");
        List<Server> servers = lb.getAllServers();
        Server serverByInstanceId = getServerByInstanceId(servers, instanceId);
        if(serverByInstanceId!=null){
            return serverByInstanceId;
        }
        //如果传的的instanceId为空或者不存在时，使用默认的轮训方式
        //return DEFAULT_RULE.choose(key);
        return super.choose(key);
    }

    private Server getServerByInstanceId(List<Server> servers, String instanceId) {
        for (Server s : servers) {
            if(s.getMetaInfo().getInstanceId().equals(instanceId)){
                return s;
            }
        }
        return null;
    }

    @Override
    public void setLoadBalancer(ILoadBalancer lb) {
        this.lb = lb;
    }
    @Override
    public ILoadBalancer getLoadBalancer() {
        return lb;
    }

}
